from kafka import KafkaProducer
from time import sleep
import requests
import pandas as pd
import os
from pathlib import Path

# Configurations
KAFKA_TOPIC = "api-gateway"  # topic名称
KAFKA_SERVERS = ["10.0.88.228:9020"]  # kafka host
ES_URL = "http://10.0.88.227:9200/api-gateway_2024-04-15/_count"  # ES查询数据条数
N_MESSAGES = 1  # 发送条数
SLEEP_TIME = 1  # 发送时长


def query_es_count():
    """查询ES中的数据总量"""
    headers = {
        'Authorization': 'Basic ZWxhc3RpYzpjSjNAcFkyKXFHMl5tTjgk'
    }
    response = requests.get(ES_URL, headers=headers)
    if response.status_code == 200:
        return response.json()['count']
    else:
        print("木有查到条数哦")
        return None


# 初始化 Kafka producer
producer = KafkaProducer(bootstrap_servers=KAFKA_SERVERS, api_version=(2, 1))

# 查询发送前的数据条数
initial_count = query_es_count()
print(f"脚本运行前es数据量: {initial_count}")

# 发送message到topic
for _ in range(N_MESSAGES):
    with open('../data/kafka10MB.txt', 'r') as message_file:
        message_content = message_file.read().encode('utf-8')
        # message_content = '123'.encode('utf-8')
        producer.send(KAFKA_TOPIC, message_content)
        sleep(SLEEP_TIME / N_MESSAGES)

# 所有信息发送完毕
producer.flush()

# 查询发送后的数据条数
final_count = query_es_count()
print(f"脚本运行后es数据量: {final_count}")

# 计算两次查询结果的差值
if initial_count is not None and final_count is not None:
    count_difference = final_count - initial_count
    print(f"es入库数据量: {count_difference}")
else:
    count_difference = None

# 记录到excel
# 检查文件是否存在，以及sheet是否存在的逻辑略去，直接处理数据写入的部分
if count_difference is not None:
    data = {
        '发送条数': [N_MESSAGES],
        '时长': [SLEEP_TIME],
        '初始条数': [initial_count],
        '最终条数': [final_count],
        '入库条数': [count_difference],
        '入库成功率': [f"{(count_difference / N_MESSAGES) * 100:.2f}%"]
    }
    df_new = pd.DataFrame(data)
    current_dir = os.path.dirname(os.path.abspath(__file__))
    filename = os.path.join(current_dir, 'kafka_es_data.xlsx')
    filepath = Path(filename)
    if filepath.exists():
        # 文件存在，尝试读取已有数据
        try:
            df_existing = pd.read_excel(filename)
            # 合并新旧DataFrame
            df_combined = pd.concat([df_existing, df_new], ignore_index=True)
        except Exception as e:  # 如果读取失败（例如文件存在但是是空的）
            print(f"Error reading the file: {e}")
            df_combined = df_new
    else:
        df_combined = df_new  # 文件不存在，直接使用新DataFrame

    # 写数据到文件，此时没有在意sheet已存在的问题，因为我们处理了整个文件的数据
    with pd.ExcelWriter(filename, mode='w', engine='openpyxl') as writer:
        df_combined.to_excel(writer, index=False)

print("脚本运行完毕.")
